"""
记忆管理器

这个模块负责管理对话历史、上下文窗口和长期记忆存储。
支持多种记忆类型，包括缓冲记忆、摘要记忆和向量记忆。
"""

import logging
import json
from typing import List, Dict, Any, Optional
from datetime import datetime, timedelta
from langchain_core.memory import BaseMemory
from langchain.memory import ConversationBufferMemory, ConversationSummaryMemory
from langchain.memory import ConversationBufferWindowMemory, ConversationSummaryBufferMemory
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
from langchain_core.chat_history import BaseChatMessageHistory
from config.settings import settings

logger = logging.getLogger(__name__)

class InMemoryChatHistory(BaseChatMessageHistory):
    """
    内存中的聊天历史实现
    
    用于存储和管理单个会话的消息历史。
    """
    
    def __init__(self):
        """初始化聊天历史"""
        self.messages: List[BaseMessage] = []
        self.created_at = datetime.now()
        self.last_accessed = datetime.now()
    
    def add_message(self, message: BaseMessage) -> None:
        """添加消息到历史记录"""
        self.messages.append(message)
        self.last_accessed = datetime.now()
    
    def clear(self) -> None:
        """清空历史记录"""
        self.messages.clear()
        self.last_accessed = datetime.now()
    
    def get_messages(self) -> List[BaseMessage]:
        """获取所有消息"""
        self.last_accessed = datetime.now()
        return self.messages.copy()

class MemoryManager:
    """
    记忆管理器类
    
    负责管理不同类型的记忆和会话历史。
    支持多会话管理、记忆持久化和自动清理。
    """
    
    def __init__(self):
        """初始化记忆管理器"""
        self.sessions: Dict[str, Dict[str, Any]] = {}
        self.default_session_id = "default"
        self.memory_types = {
            "buffer": ConversationBufferMemory,
            "window": ConversationBufferWindowMemory,
            "summary": ConversationSummaryMemory,
            "summary_buffer": ConversationSummaryBufferMemory
        }
        
        # 创建默认会话
        self.create_session(self.default_session_id)
        
        logger.info("记忆管理器初始化完成")
    
    def create_session(self, session_id: str, memory_type: str = "buffer", **kwargs) -> bool:
        """
        创建新的会话
        
        Args:
            session_id (str): 会话ID
            memory_type (str): 记忆类型 (buffer, window, summary, summary_buffer)
            **kwargs: 记忆配置参数
            
        Returns:
            bool: 是否创建成功
        """
        if session_id in self.sessions:
            logger.warning(f"会话 {session_id} 已存在")
            return False
        
        try:
            # 创建聊天历史
            chat_history = InMemoryChatHistory()
            
            # 配置记忆参数
            memory_config = {
                "chat_memory": chat_history,
                "return_messages": True
            }
            
            # 根据记忆类型添加特定配置
            if memory_type == "window":
                memory_config["k"] = kwargs.get("window_size", settings.MAX_CONVERSATION_HISTORY)
            elif memory_type in ["summary", "summary_buffer"]:
                from core.llm_manager import llm_manager
                memory_config["llm"] = llm_manager.get_current_model()
                if memory_type == "summary_buffer":
                    memory_config["max_token_limit"] = kwargs.get("max_tokens", 1000)
            
            # 创建记忆实例
            memory_class = self.memory_types[memory_type]
            memory = memory_class(**memory_config)
            
            # 存储会话信息
            self.sessions[session_id] = {
                "memory": memory,
                "chat_history": chat_history,
                "memory_type": memory_type,
                "created_at": datetime.now(),
                "last_accessed": datetime.now(),
                "metadata": kwargs.get("metadata", {})
            }
            
            logger.info(f"已创建会话 {session_id}，记忆类型: {memory_type}")
            return True
            
        except Exception as e:
            logger.error(f"创建会话 {session_id} 失败: {str(e)}")
            return False
    
    def get_session(self, session_id: str = None) -> Optional[Dict[str, Any]]:
        """
        获取会话信息
        
        Args:
            session_id (str, optional): 会话ID，默认使用默认会话
            
        Returns:
            Optional[Dict[str, Any]]: 会话信息
        """
        if session_id is None:
            session_id = self.default_session_id
        
        if session_id not in self.sessions:
            logger.warning(f"会话 {session_id} 不存在")
            return None
        
        session = self.sessions[session_id]
        session["last_accessed"] = datetime.now()
        return session
    
    def add_message(self, message: BaseMessage, session_id: str = None) -> bool:
        """
        添加消息到会话
        
        Args:
            message (BaseMessage): 要添加的消息
            session_id (str, optional): 会话ID
            
        Returns:
            bool: 是否添加成功
        """
        session = self.get_session(session_id)
        if not session:
            return False
        
        try:
            session["chat_history"].add_message(message)
            logger.debug(f"已添加消息到会话 {session_id or self.default_session_id}")
            return True
            
        except Exception as e:
            logger.error(f"添加消息失败: {str(e)}")
            return False
    
    def add_user_message(self, content: str, session_id: str = None) -> bool:
        """
        添加用户消息
        
        Args:
            content (str): 消息内容
            session_id (str, optional): 会话ID
            
        Returns:
            bool: 是否添加成功
        """
        message = HumanMessage(content=content)
        return self.add_message(message, session_id)
    
    def add_ai_message(self, content: str, session_id: str = None) -> bool:
        """
        添加AI消息
        
        Args:
            content (str): 消息内容
            session_id (str, optional): 会话ID
            
        Returns:
            bool: 是否添加成功
        """
        message = AIMessage(content=content)
        return self.add_message(message, session_id)
    
    def get_memory(self, session_id: str = None) -> Optional[BaseMemory]:
        """
        获取会话的记忆对象
        
        Args:
            session_id (str, optional): 会话ID
            
        Returns:
            Optional[BaseMemory]: 记忆对象
        """
        session = self.get_session(session_id)
        if session:
            return session["memory"]
        return None
    
    def get_chat_history(self, session_id: str = None) -> List[BaseMessage]:
        """
        获取聊天历史
        
        Args:
            session_id (str, optional): 会话ID
            
        Returns:
            List[BaseMessage]: 消息列表
        """
        session = self.get_session(session_id)
        if session:
            return session["chat_history"].get_messages()
        return []
    
    def get_conversation_string(self, session_id: str = None) -> str:
        """
        获取对话的字符串表示
        
        Args:
            session_id (str, optional): 会话ID
            
        Returns:
            str: 对话字符串
        """
        messages = self.get_chat_history(session_id)
        conversation_parts = []
        
        for message in messages:
            if isinstance(message, HumanMessage):
                conversation_parts.append(f"用户: {message.content}")
            elif isinstance(message, AIMessage):
                conversation_parts.append(f"助手: {message.content}")
            elif isinstance(message, SystemMessage):
                conversation_parts.append(f"系统: {message.content}")
        
        return "\n".join(conversation_parts)
    
    def clear_session(self, session_id: str = None) -> bool:
        """
        清空会话历史
        
        Args:
            session_id (str, optional): 会话ID
            
        Returns:
            bool: 是否清空成功
        """
        session = self.get_session(session_id)
        if not session:
            return False
        
        try:
            session["chat_history"].clear()
            session["memory"].clear()
            logger.info(f"已清空会话 {session_id or self.default_session_id}")
            return True
            
        except Exception as e:
            logger.error(f"清空会话失败: {str(e)}")
            return False
    
    def delete_session(self, session_id: str) -> bool:
        """
        删除会话
        
        Args:
            session_id (str): 会话ID
            
        Returns:
            bool: 是否删除成功
        """
        if session_id == self.default_session_id:
            logger.warning("不能删除默认会话")
            return False
        
        if session_id not in self.sessions:
            logger.warning(f"会话 {session_id} 不存在")
            return False
        
        del self.sessions[session_id]
        logger.info(f"已删除会话 {session_id}")
        return True
    
    def list_sessions(self) -> Dict[str, Dict[str, Any]]:
        """
        列出所有会话
        
        Returns:
            Dict[str, Dict[str, Any]]: 会话信息字典
        """
        session_info = {}
        
        for session_id, session in self.sessions.items():
            session_info[session_id] = {
                "memory_type": session["memory_type"],
                "created_at": session["created_at"].isoformat(),
                "last_accessed": session["last_accessed"].isoformat(),
                "message_count": len(session["chat_history"].get_messages()),
                "metadata": session["metadata"]
            }
        
        return session_info
    
    def cleanup_expired_sessions(self, max_age_hours: int = 24) -> int:
        """
        清理过期的会话
        
        Args:
            max_age_hours (int): 最大存活时间（小时）
            
        Returns:
            int: 清理的会话数量
        """
        cutoff_time = datetime.now() - timedelta(hours=max_age_hours)
        expired_sessions = []
        
        for session_id, session in self.sessions.items():
            if (session_id != self.default_session_id and 
                session["last_accessed"] < cutoff_time):
                expired_sessions.append(session_id)
        
        for session_id in expired_sessions:
            self.delete_session(session_id)
        
        if expired_sessions:
            logger.info(f"已清理 {len(expired_sessions)} 个过期会话")
        
        return len(expired_sessions)
    
    def export_session(self, session_id: str = None) -> Optional[Dict[str, Any]]:
        """
        导出会话数据
        
        Args:
            session_id (str, optional): 会话ID
            
        Returns:
            Optional[Dict[str, Any]]: 会话数据
        """
        session = self.get_session(session_id)
        if not session:
            return None
        
        messages = session["chat_history"].get_messages()
        exported_messages = []
        
        for message in messages:
            exported_messages.append({
                "type": message.__class__.__name__,
                "content": message.content,
                "timestamp": getattr(message, "timestamp", None)
            })
        
        return {
            "session_id": session_id or self.default_session_id,
            "memory_type": session["memory_type"],
            "created_at": session["created_at"].isoformat(),
            "messages": exported_messages,
            "metadata": session["metadata"]
        }
    
    def get_memory_stats(self) -> Dict[str, Any]:
        """
        获取记忆统计信息
        
        Returns:
            Dict[str, Any]: 统计信息
        """
        total_sessions = len(self.sessions)
        total_messages = sum(
            len(session["chat_history"].get_messages()) 
            for session in self.sessions.values()
        )
        
        memory_type_counts = {}
        for session in self.sessions.values():
            memory_type = session["memory_type"]
            memory_type_counts[memory_type] = memory_type_counts.get(memory_type, 0) + 1
        
        return {
            "total_sessions": total_sessions,
            "total_messages": total_messages,
            "memory_type_distribution": memory_type_counts,
            "default_session_id": self.default_session_id
        }

# 创建全局记忆管理器实例
memory_manager = MemoryManager() 